#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
数据清理脚本
检查和清理bearing_features.csv文件中的数据类型问题
"""

import pandas as pd
import numpy as np
import os

def clean_bearing_features(input_file='bearing_features.csv', output_file='bearing_features_cleaned.csv'):
    """清理bearing_features.csv文件"""
    print("=== 数据清理脚本 ===")
    
    if not os.path.exists(input_file):
        print(f"错误: 文件不存在 {input_file}")
        return None
    
    # 读取数据
    print(f"读取数据: {input_file}")
    df = pd.read_csv(input_file)
    print(f"原始数据形状: {df.shape}")
    print(f"原始列数: {len(df.columns)}")
    
    # 显示列信息
    print("\n=== 列信息 ===")
    for i, col in enumerate(df.columns):
        dtype = df[col].dtype
        unique_count = df[col].nunique()
        null_count = df[col].isnull().sum()
        print(f"{i+1:3d}. {col:25s} | {str(dtype):15s} | 唯一值: {unique_count:3d} | 空值: {null_count:3d}")
    
    # 识别非数值列
    print("\n=== 非数值列分析 ===")
    non_numeric_cols = []
    for col in df.columns:
        if not pd.api.types.is_numeric_dtype(df[col]):
            non_numeric_cols.append(col)
            print(f"非数值列: {col} (类型: {df[col].dtype})")
            if df[col].nunique() < 20:  # 如果唯一值较少，显示所有值
                print(f"  唯一值: {df[col].unique()}")
    
    # 定义要排除的列
    exclude_cols = ['label', 'filename', 'rpm', 'sampling_rate', 'bearing_type', 'signal_length']
    exclude_cols.extend(non_numeric_cols)
    exclude_cols = list(set(exclude_cols))  # 去重
    
    print(f"\n将排除的列: {exclude_cols}")
    
    # 选择数值特征
    numeric_cols = []
    for col in df.columns:
        if col not in exclude_cols and pd.api.types.is_numeric_dtype(df[col]):
            numeric_cols.append(col)
    
    print(f"\n数值特征列数: {len(numeric_cols)}")
    
    # 创建清理后的数据
    cleaned_df = df.copy()
    
    # 处理数值列
    for col in numeric_cols:
        # 转换为数值类型
        cleaned_df[col] = pd.to_numeric(cleaned_df[col], errors='coerce')
        
        # 填充NaN值
        cleaned_df[col] = cleaned_df[col].fillna(0)
        
        # 处理无穷值
        cleaned_df[col] = cleaned_df[col].replace([np.inf, -np.inf], 0)
    
    # 保留必要的元信息列
    meta_cols = ['label', 'filename', 'rpm', 'sampling_rate']
    for col in meta_cols:
        if col in df.columns:
            cleaned_df[col] = df[col]
    
    # 保存清理后的数据
    cleaned_df.to_csv(output_file, index=False, encoding='utf-8-sig')
    print(f"\n清理后的数据已保存到: {output_file}")
    print(f"清理后数据形状: {cleaned_df.shape}")
    
    # 显示清理后的列信息
    print("\n=== 清理后数值列信息 ===")
    for col in numeric_cols[:10]:  # 只显示前10列
        mean_val = cleaned_df[col].mean()
        std_val = cleaned_df[col].std()
        min_val = cleaned_df[col].min()
        max_val = cleaned_df[col].max()
        print(f"{col:25s} | 均值: {mean_val:8.4f} | 标准差: {std_val:8.4f} | 范围: [{min_val:8.4f}, {max_val:8.4f}]")
    
    if len(numeric_cols) > 10:
        print(f"... 还有 {len(numeric_cols) - 10} 个数值列")
    
    return cleaned_df, numeric_cols

def main():
    """主函数"""
    print("=== 轴承特征数据清理工具 ===")
    
    # 清理数据
    cleaned_df, numeric_cols = clean_bearing_features()
    
    if cleaned_df is not None:
        print(f"\n=== 清理完成 ===")
        print(f"数值特征数量: {len(numeric_cols)}")
        print(f"建议使用清理后的文件: bearing_features_cleaned.csv")
        print("\n现在可以运行:")
        print("python ablation_study.py")
        print("python pca_feature_reduction.py")
        print("python tree_model_shap_analysis.py")

if __name__ == "__main__":
    main()

